# --- STEP 1: Import libraries ---
import requests
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
import json
import logging

# --- STEP 2: Set up logging ---
logging.basicConfig(
    level=logging.INFO,
    format='%(message)s',
    handlers=[
        logging.FileHandler(f'qwen_api_communication_multiple_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'),
        logging.StreamHandler()
    ]
)

# --- STEP 3: Set your Qwen API Key ---
QWEN_API_KEY = "sk-6b3fb2c924084418b1332f6161f4be2a"  # Replace with your Qwen API key

# --- STEP 4: Define core functions ---
# 📍 Get coordinates for a city
def get_coordinates(city):
    logging.info(f"📌 [get_coordinates] Looking up coordinates for: {city}")
    response = requests.get(
        "https://geocoding-api.open-meteo.com/v1/search",
        params={"name": city, "count": 1, "language": "en"}
    )
    data = response.json()
    if data.get("results"):
        loc = data["results"][0]
        logging.info(f"✅ Found: {loc['name']} → (lat: {loc['latitude']}, lon: {loc['longitude']})")
        return {"latitude": loc["latitude"], "longitude": loc["longitude"], "name": loc["name"]}
    else:
        logging.info("❌ City not found.")
        return None

# 🌡️ Download daily temperature data for a year
def get_temperature_data(lat, lon, year):
    logging.info(f"🌡️ [get_temperature_data] Fetching daily max temperature for {year} at (lat: {lat}, lon: {lon})")
    url = "https://archive-api.open-meteo.com/v1/archive"
    params = {
        "latitude": lat,
        "longitude": lon,
        "start_date": f"{year}-01-01",
        "end_date": f"{year}-12-31",
        "daily": "temperature_2m_max",
        "timezone": "auto"
    }
    r = requests.get(url, params=params)
    data = r.json()
    logging.info(f"✅ Retrieved {len(data['daily']['time'])} days of data.")
    return pd.DataFrame({
        "date": pd.to_datetime(data["daily"]["time"]),
        "temperature": data["daily"]["temperature_2m_max"]
    })

# 🔄 Compare temperatures between two cities
def compare_cities_temperature(city1, city2, year):
    logging.info(f"\n🔄 [compare_cities_temperature] Comparing {city1} and {city2} in {year}")
    
    # Get coordinates for both cities
    coords1 = get_coordinates(city1)
    coords2 = get_coordinates(city2)
    
    if not coords1 or not coords2:
        return "❌ Could not find coordinates for one or both cities."
    
    # Get temperature data for both cities
    df1 = get_temperature_data(coords1["latitude"], coords1["longitude"], year)
    df2 = get_temperature_data(coords2["latitude"], coords2["longitude"], year)
    
    # Calculate statistics
    stats1 = {
        "max": df1['temperature'].max(),
        "min": df1['temperature'].min(),
        "avg": df1['temperature'].mean().round(1)
    }
    stats2 = {
        "max": df2['temperature'].max(),
        "min": df2['temperature'].min(),
        "avg": df2['temperature'].mean().round(1)
    }
    
    # Plot comparison chart
    plt.figure(figsize=(14, 6))
    plt.plot(df1['date'], df1['temperature'], color='red', linewidth=1.5, label=city1)
    plt.plot(df2['date'], df2['temperature'], color='blue', linewidth=1.5, label=city2)
    plt.title(f"Temperature Comparison: {city1} vs {city2} in {year}")
    plt.xlabel("Date")
    plt.ylabel("Temperature (°C)")
    plt.grid(True, alpha=0.3, linestyle="--")
    plt.legend()
    plt.tight_layout()
    plt.show()
    
    # Generate comparison summary
    summary = (
        f"📊 Temperature Comparison between {city1} and {city2} in {year}:\n\n"
        f"{city1}:\n"
        f"🔹 Max Temp: {stats1['max']}°C\n"
        f"🔹 Min Temp: {stats1['min']}°C\n"
        f"🔹 Avg Temp: {stats1['avg']}°C\n\n"
        f"{city2}:\n"
        f"🔹 Max Temp: {stats2['max']}°C\n"
        f"🔹 Min Temp: {stats2['min']}°C\n"
        f"🔹 Avg Temp: {stats2['avg']}°C\n\n"
        f"📈 A comparison chart is shown above."
    )
    
    return summary

# --- STEP 5: Define function calling with Qwen API ---
def ask_qwen(user_question):
    url = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {QWEN_API_KEY}",
        "Content-Type": "application/json"
    }
    
    # Function definitions for Qwen
    functions = [
        {
            "name": "get_coordinates",
            "description": "Get latitude and longitude for a city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "City name"}
                },
                "required": ["city"]
            }
        },
        {
            "name": "get_temperature_data",
            "description": "Get temperature data for a location and year.",
            "parameters": {
                "type": "object",
                "properties": {
                    "lat": {"type": "number", "description": "Latitude"},
                    "lon": {"type": "number", "description": "Longitude"},
                    "year": {"type": "integer", "description": "Year (e.g., 2023)"}
                },
                "required": ["lat", "lon", "year"]
            }
        },
        {
            "name": "compare_cities_temperature",
            "description": "Compare temperature data between two cities for a given year.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city1": {"type": "string", "description": "First city name"},
                    "city2": {"type": "string", "description": "Second city name"},
                    "year": {"type": "integer", "description": "Year (e.g., 2023)"}
                },
                "required": ["city1", "city2", "year"]
            }
        }
    ]
    
    # Initialize messages with system and user content
    messages = [
        {
            "role": "system",
            "content": "You are a helpful assistant that can get weather information and create temperature charts. You can compare temperatures between cities. When asked to list coordinates of multiple cities, you should make separate function calls for each city."
        },
        {
            "role": "user",
            "content": user_question
        }
    ]
    
    try:
        # Log the start of the query
        logging.info("\n" + "="*50)
        logging.info("=== QUERY START ===")
        logging.info("="*50)
        logging.info(f"User Question: {user_question}")
        logging.info("="*50)
        
        logging.info("\n" + "="*50)
        logging.info("=== Local to Qwen API Request ===")
        logging.info("="*50)
        logging.info(f"Request URL: {url}")
        logging.info(f"Request Headers: {json.dumps(headers, indent=2)}")
        
        # Continue making API calls until there are no more function calls
        function_call_count = 0
        while True:
            data = {
                "model": "qwen-turbo",
                "messages": messages,
                "functions": functions,
                "function_call": "auto"
            }
            
            logging.info(f"Request Body: {json.dumps(data, indent=2)}")
            
            response = requests.post(url, headers=headers, json=data)
            response.raise_for_status()
            message = response.json()["choices"][0]["message"]
            
            logging.info("\n" + "="*50)
            logging.info("=== Qwen API Response ===")
            logging.info("="*50)
            logging.info(f"Response: {json.dumps(message, indent=2)}")
            
            # Log the LLM's reasoning process
            if "content" in message and message["content"]:
                logging.info("\n=== LLM Reasoning Process ===")
                logging.info(message["content"])
            
            logging.info("="*50)
            
            # Add the assistant's message to the conversation
            messages.append(message)
            
            # If there's no function call, we're done
            if "function_call" not in message:
                logging.info("\n=== Function Calling Process Complete ===")
                logging.info("Reason: No more function calls in LLM response")
                print("\n🧠 Qwen says:")
                print(message["content"])
                break
                
            # Process the function call
            function_call_count += 1
            fn_name = message["function_call"]["name"]
            args = json.loads(message["function_call"]["arguments"])
            
            # Log the start of the function call
            logging.info("\n" + "="*50)
            logging.info(f"=== FUNCTION CALL #{function_call_count} START: {fn_name} ===")
            logging.info("="*50)
            logging.info(f"Arguments: {json.dumps(args, indent=2)}")
            
            # Call the matching Python function
            if fn_name == "get_coordinates":
                coords = get_coordinates(args["city"])
                if coords:
                    result = f"{args['city'].title()} is located at (lat: {coords['latitude']}, lon: {coords['longitude']})."
                else:
                    result = f"Could not find coordinates for {args['city']}."
                
                # Add the function result to messages
                messages.append({"role": "function", "name": fn_name, "content": result})
                
            elif fn_name == "get_temperature_data":
                df = get_temperature_data(args["lat"], args["lon"], args["year"])
                result = f"Retrieved {len(df)} days of temperature data."
                messages.append({"role": "function", "name": fn_name, "content": result})
                
            elif fn_name == "compare_cities_temperature":
                result = compare_cities_temperature(args["city1"], args["city2"], args["year"])
                messages.append({"role": "function", "name": fn_name, "content": result})
                
            else:
                logging.info("\n=== Function Calling Process Complete ===")
                logging.info(f"Reason: Unknown function '{fn_name}' called")
                print("\n❌ Unknown function called.")
                break
            
            # Log the end of the function call
            logging.info(f"Function Result: {result}")
            logging.info("="*50)
            logging.info(f"=== FUNCTION CALL #{function_call_count} END: {fn_name} ===")
            logging.info("="*50)
            
        # Log the end of the query
        logging.info("\n" + "="*50)
        logging.info("=== QUERY END ===")
        logging.info("="*50)
            
    except requests.exceptions.RequestException as e:
        logging.error(f"API Error: {str(e)}")
    except json.JSONDecodeError as e:
        logging.error(f"JSON Parse Error: {str(e)}")
    except Exception as e:
        logging.error(f"Unexpected Error: {str(e)}")

# --- STEP 6: Example usage ---
if __name__ == "__main__":
    # Example 1: Compare temperatures between two cities
    # ask_qwen("Compare the temperature trends between Beijing and Shanghai in 2023")
    ask_qwen("list the coordinates of Beijing, Kunming, Fuzhou and Shanghai. And Compare the temperature trends between cities have lowest latitude in 2023")
